# -*- coding: utf-8 -*-
# imageio is distributed under the terms of the (new) BSD License.

""" SPE file reader
"""

from collections import namedtuple
from datetime import datetime
import logging
import os
from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Union

import numpy as np

from .. import formats
from ..core import Format


logger = logging.getLogger(__name__)


class Spec:
    """SPE file specification data

    Tuples of (offset, datatype, count), where offset is the offset in the SPE
    file and datatype is the datatype as used in `numpy.fromfile`()

    `data_start` is the offset of actual image data.

    `dtypes` translates SPE datatypes (0...4) to numpy ones, e. g. dtypes[0]
    is dtype("<f") (which is np.float32).

    `controllers` maps the `type` metadata to a human readable name

    `readout_modes` maps the `readoutMode` metadata to something human readable
    although this may not be accurate since there is next to no documentation
    to be found.
    """

    basic = {
        "datatype": (108, "<h"),  # dtypes
        "xdim": (42, "<H"),
        "ydim": (656, "<H"),
        "xml_footer_offset": (678, "<Q"),
        "NumFrames": (1446, "<i"),
        "file_header_ver": (1992, "<f"),
    }

    metadata = {
        # ROI information
        "NumROI": (1510, "<h"),
        "ROIs": (
            1512,
            np.dtype(
                [
                    ("startx", "<H"),
                    ("endx", "<H"),
                    ("groupx", "<H"),
                    ("starty", "<H"),
                    ("endy", "<H"),
                    ("groupy", "<H"),
                ]
            ),
            10,
        ),
        # chip-related sizes
        "xDimDet": (6, "<H"),
        "yDimDet": (18, "<H"),
        "VChipXdim": (14, "<h"),
        "VChipYdim": (16, "<h"),
        # other stuff
        "controller_version": (0, "<h"),
        "logic_output": (2, "<h"),
        "amp_high_cap_low_noise": (4, "<H"),  # enum?
        "mode": (8, "<h"),  # enum?
        "exposure_sec": (10, "<f"),
        "date": (20, "<10S"),
        "detector_temp": (36, "<f"),
        "detector_type": (40, "<h"),
        "st_diode": (44, "<h"),
        "delay_time": (46, "<f"),
        # shutter_control: normal, disabled open, disabled closed
        # But which one is which?
        "shutter_control": (50, "<H"),
        "absorb_live": (52, "<h"),
        "absorb_mode": (54, "<H"),
        "can_do_virtual_chip": (56, "<h"),
        "threshold_min_live": (58, "<h"),
        "threshold_min_val": (60, "<f"),
        "threshold_max_live": (64, "<h"),
        "threshold_max_val": (66, "<f"),
        "time_local": (172, "<7S"),
        "time_utc": (179, "<7S"),
        "adc_offset": (188, "<H"),
        "adc_rate": (190, "<H"),
        "adc_type": (192, "<H"),
        "adc_resolution": (194, "<H"),
        "adc_bit_adjust": (196, "<H"),
        "gain": (198, "<H"),
        "comments": (200, "<80S", 5),
        "geometric": (600, "<H"),  # flags
        "sw_version": (688, "<16S"),
        "spare_4": (742, "<436S"),
        "XPrePixels": (98, "<h"),
        "XPostPixels": (100, "<h"),
        "YPrePixels": (102, "<h"),
        "YPostPixels": (104, "<h"),
        "readout_time": (672, "<f"),
        "xml_footer_offset": (678, "<Q"),
        "type": (704, "<h"),  # controllers
        "clockspeed_us": (1428, "<f"),
        "readout_mode": (1480, "<H"),  # readout_modes
        "window_size": (1482, "<H"),
        "file_header_ver": (1992, "<f"),
    }

    data_start = 4100

    dtypes = {
        0: np.dtype(np.float32),
        1: np.dtype(np.int32),
        2: np.dtype(np.int16),
        3: np.dtype(np.uint16),
        8: np.dtype(np.uint32),
    }

    controllers = [
        "new120 (Type II)",
        "old120 (Type I)",
        "ST130",
        "ST121",
        "ST138",
        "DC131 (PentaMax)",
        "ST133 (MicroMax/Roper)",
        "ST135 (GPIB)",
        "VTCCD",
        "ST116 (GPIB)",
        "OMA3 (GPIB)",
        "OMA4",
    ]

    # This was gathered from random places on the internet and own experiments
    # with the camera. May not be accurate.
    readout_modes = ["full frame", "frame transfer", "kinetics"]

    # Do not decode the following metadata keys into strings, but leave them
    # as byte arrays
    no_decode = ["spare_4"]


class SDTControlSpec:
    """Extract metadata written by the SDT-control software

    Some of it is encoded in the comment strings
    (see :py:meth:`parse_comments`). Also, date and time are encoded in a
    peculiar way (see :py:meth:`get_datetime`). Use :py:meth:`extract_metadata`
    to update the metadata dict.
    """
    months = {
        # Convert SDT-control month strings to month numbers
        "Jän": 1, "Jan": 1, "Feb": 2, "Mär": 3, "Mar": 3, "Apr": 4, "Mai": 5,
        "May": 5, "Jun": 6, "Jul": 7, "Aug": 8, "Sep": 9, "Okt": 10, "Oct": 10,
        "Nov": 11, "Dez": 12, "Dec": 12
    }

    sequence_types = {
        # TODO: complete
        "SEQU": "standard", "SETO": "TOCCSL", "KINE": "kinetics",
        "SEAR": "arbitrary"
    }

    class CommentDesc:
        """Describe how to extract a metadata entry from a comment string"""
        n: int
        """Which of the 5 SPE comment fields to use."""
        slice: slice
        """Which characters from the `n`-th comment to use."""
        cvt: Callable[[str], Any]
        """How to convert characters to something useful."""
        scale: Union[None, float]
        """Optional scaling factor for numbers"""
        def __init__(self, n: int, slice: slice,
                     cvt: Callable[[str], Any] = str,
                     scale: Optional[float] = None):
            self.n = n
            self.slice = slice
            self.cvt = cvt
            self.scale = scale

    comments = {
        "sdt_major_version": CommentDesc(4, slice(66, 68), int),
        "sdt_minor_version": CommentDesc(4, slice(68, 70), int),
        "sdt_controller_name": CommentDesc(4, slice(0, 6), str),
        "exposure_time": CommentDesc(1, slice(64, 73), float, 10**-6),
        "color_code": CommentDesc(4, slice(10, 14), str),
        "detection_channels": CommentDesc(4, slice(15, 16), int),
        "background_subtraction": CommentDesc(4, 14, lambda x: x == "B"),
        "em_active": CommentDesc(4, 32, lambda x: x == "E"),
        "em_gain": CommentDesc(4, slice(28, 32), int),
        "modulation_active": CommentDesc(4, 33, lambda x: x == "A"),
        "pixel_size": CommentDesc(4, slice(25, 28), float, 0.1),
        "sequence_type": CommentDesc(
            4, slice(6, 10), lambda x: __class__.sequence_types[x]),
        "grid": CommentDesc(4, slice(16, 25), float, 10**-6),
        "n_macro": CommentDesc(1, slice(0, 4), int),
        "delay_macro": CommentDesc(1, slice(10, 19), float, 10**-3),
        "n_mini": CommentDesc(1, slice(4, 7), int),
        "delay_mini": CommentDesc(1, slice(19, 28), float, 10**-6),
        "n_micro": CommentDesc(1, slice(7, 10), int),
        "delay_micro": CommentDesc(1, slice(28, 37), float, 10**-6),
        "n_subpics": CommentDesc(1, slice(7, 10), int),
        "delay_shutter": CommentDesc(1, slice(73, 79), float, 10**-6),
        "delay_prebleach": CommentDesc(1, slice(37, 46), float, 10**-6),
        "bleach_time": CommentDesc(1, slice(46, 55), float, 10**-6),
        "recovery_time": CommentDesc(1, slice(55, 64), float, 10**-6)
    }

    @staticmethod
    def parse_comments(comments: Sequence[str]) -> Union[Dict, None]:
        """Extract SDT-control metadata from comments

        Parameters
        ----------
        comments
            List of SPE file comments, typically ``metadata["comments"]``.

        Returns
        -------
        If SDT-control comments were detected, return a dict of metadata, else
        `None`.
        """
        sdt_md = {}
        if comments[4][70:] != "COMVER0500":
            logger.debug("SDT-control comments not found.")
            return None

        sdt_md = {}
        for name, spec in SDTControlSpec.comments.items():
            try:
                v = spec.cvt(comments[spec.n][spec.slice])
                if spec.scale is not None:
                    v *= spec.scale
            except Exception as e:
                logger.debug("Failed to decode SDT-control metadata "
                             f'field "{name}": {e}')
            sdt_md[name] = v
        comment = comments[0] + comments[2]
        sdt_md["comment"] = comment.strip()
        return sdt_md

    @staticmethod
    def get_datetime(date: str, time: str) -> Union[datetime, None]:
        """Turn date and time saved by SDT-control into proper datetime object

        Parameters
        ----------
        date
            SPE file date, typically ``metadata["date"]``.
        time
            SPE file date, typically ``metadata["time_local"]``.

        Returns
        -------
        File's datetime if parsing was succsessful, else None.
        """
        try:
            month = __class__.months[date[2:5]]
            return datetime(
                int(date[5:9]), month, int(date[0:2]), int(time[0:2]),
                int(time[2:4]), int(time[4:6]))
        except Exception as e:
            logger.info(
                f"Failed to decode date from SDT-control metadata: {e}.")

    @staticmethod
    def extract_metadata(meta: Mapping, char_encoding: str = "latin1"):
        """Extract SDT-control metadata from SPE metadata

        SDT-control stores some metadata in comments and other fields.
        Extract them and remove unused entries.

        Parameters
        ----------
        meta
            SPE file metadata. Modified in place.
        char_encoding
            Character encoding used to decode strings in the metadata.
        """
        sdt_meta = __class__.parse_comments(meta["comments"])
        if not sdt_meta:
            return
        # This file has SDT-control metadata
        meta.pop("comments")
        meta.update(sdt_meta)

        # Get date and time in a usable format
        dt = __class__.get_datetime(
            meta["date"], meta["time_local"])
        if dt:
            meta["datetime"] = dt
            meta.pop("date")
            meta.pop("time_local")

        sp4 = meta["spare_4"]
        try:
            meta["modulation_script"] = sp4.decode(char_encoding)
            meta.pop("spare_4")
        except UnicodeDecodeError:
            logger.warning("Failed to decode SDT-control laser "
                           "modulation script. Bad char_encoding?")

        # Get rid of unused data
        meta.pop("time_utc")
        meta.pop("exposure_sec")


class SpeFormat(Format):
    """Some CCD camera software produces images in the Princeton Instruments
    SPE file format. This plugin supports reading such files.

    Parameters for reading
    ----------------------
    char_encoding : str
        Character encoding used to decode strings in the metadata. Defaults
        to "latin1".
    check_filesize : bool
        The number of frames in the file is stored in the file header. However,
        this number may be wrong for certain software. If this is `True`
        (default), derive the number of frames also from the file size and
        raise a warning if the two values do not match.
    sdt_meta : bool
        If set to `True` (default), check for special metadata written by the
        `SDT-control` software. Does not have an effect for files written by
        other software.

    Metadata for reading
    --------------------
    ROIs : list of dict
        Regions of interest used for recording images. Each dict has the
        "top_left" key containing x and y coordinates of the top left corner,
        the "bottom_right" key with x and y coordinates of the bottom right
        corner, and the "bin" key with number of binned pixels in x and y
        directions.
    comments : list of str
        The SPE format allows for 5 comment strings of 80 characters each.
    controller_version : int
        Hardware version
    logic_output : int
        Definition of output BNC
    amp_hi_cap_low_noise : int
        Amp switching mode
    mode : int
        Timing mode
    exp_sec : float
        Alternative exposure in seconds
    date : str
        Date string
    detector_temp : float
        Detector temperature
    detector_type : int
        CCD / diode array type
    st_diode : int
        Trigger diode
    delay_time : float
        Used with async mode
    shutter_control : int
        Normal, disabled open, or disabled closed
    absorb_live : bool
        on / off
    absorb_mode : int
        Reference strip or file
    can_do_virtual_chip : bool
        True or False whether chip can do virtual chip
    threshold_min_live : bool
        on / off
    threshold_min_val : float
        Threshold minimum value
    threshold_max_live : bool
        on / off
    threshold_max_val : float
        Threshold maximum value
    time_local : str
        Experiment local time
    time_utc : str
        Experiment UTC time
    adc_offset : int
        ADC offset
    adc_rate : int
        ADC rate
    adc_type : int
        ADC type
    adc_resolution : int
        ADC resolution
    adc_bit_adjust : int
        ADC bit adjust
    gain : int
        gain
    sw_version : str
        Version of software which created this file
    spare_4 : bytes
        Reserved space
    readout_time : float
        Experiment readout time
    type : str
        Controller type
    clockspeed_us : float
        Vertical clock speed in microseconds
    readout_mode : {"full frame", "frame transfer", "kinetics", ""}
        Readout mode. Empty string means that this was not set by the
        Software.
    window_size : int
        Window size for Kinetics mode
    file_header_ver : float
        File header version
    chip_size : [int, int]
        x and y dimensions of the camera chip
    virt_chip_size : [int, int]
        Virtual chip x and y dimensions
    pre_pixels : [int, int]
        Pre pixels in x and y dimensions
    post_pixels : [int, int],
        Post pixels in x and y dimensions
    geometric : list of {"rotate", "reverse", "flip"}
        Geometric operations
    sdt_major_version : int (only for files created by SDT-control)
        Major version of SDT-control software
    sdt_minor_version : int (only for files created by SDT-control)
        Minor version of SDT-control software
    sdt_controller_name : str (only for files created by SDT-control)
        Controller name
    exposure_time : float (only for files created by SDT-control)
        Exposure time in seconds
    color_code : str (only for files created by SDT-control)
        Color channels used
    detection_channels : int (only for files created by SDT-control)
        Number of channels
    background_subtraction : bool (only for files created by SDT-control)
        Whether background subtraction war turned on
    em_active : bool (only for files created by SDT-control)
        Whether EM was turned on
    em_gain : int (only for files created by SDT-control)
        EM gain
    modulation_active : bool (only for files created by SDT-control)
        Whether laser modulation (“attenuate”) was turned on
    pixel_size : float (only for files created by SDT-control)
        Camera pixel size
    sequence_type : str (only for files created by SDT-control)
        Type of sequnce (standard, TOCCSL, arbitrary, …)
    grid : float (only for files created by SDT-control)
        Sequence time unit (“grid size”) in seconds
    n_macro : int (only for files created by SDT-control)
        Number of macro loops
    delay_macro : float (only for files created by SDT-control)
        Time between macro loops in seconds
    n_mini : int (only for files created by SDT-control)
        Number of mini loops
    delay_mini : float (only for files created by SDT-control)
        Time between mini loops in seconds
    n_micro : int (only for files created by SDT-control)
        Number of micro loops
    delay_micro : float (only for files created by SDT-control)
        Time between micro loops in seconds
    n_subpics : int (only for files created by SDT-control)
        Number of sub-pictures
    delay_shutter : float (only for files created by SDT-control)
        Camera shutter delay in seconds
    delay_prebleach : float (only for files created by SDT-control)
        Pre-bleach delay in seconds
    bleach_time : float (only for files created by SDT-control)
        Bleaching time in seconds
    recovery_time : float (only for files created by SDT-control)
        Recovery time in seconds
    comment : str (only for files created by SDT-control)
        User-entered comment. This replaces the "comments" field.
    datetime : datetime.datetime (only for files created by SDT-control)
        Combines the "date" and "time_local" keys. The latter two plus
        "time_utc" are removed.
    modulation_script : str (only for files created by SDT-control)
        Laser modulation script. Replaces the "spare_4" key.
     """

    def _can_read(self, request):
        return (
            request.mode[1] in self.modes + "?" and request.extension in self.extensions
        )

    def _can_write(self, request):
        return False

    class Reader(Format.Reader):
        def _open(self, char_encoding="latin1", check_filesize=True,
                  sdt_meta=True):
            self._file = self.request.get_file()
            self._char_encoding = char_encoding

            info = self._parse_header(Spec.basic)
            self._file_header_ver = info["file_header_ver"]
            self._dtype = Spec.dtypes[info["datatype"]]
            self._shape = (info["ydim"], info["xdim"])
            self._len = info["NumFrames"]
            self._sdt_meta = sdt_meta

            if check_filesize:
                # Some software writes incorrect `NumFrames` metadata.
                # To determine the number of frames, check the size of the data
                # segment -- until the end of the file for SPE<3, until the
                # xml footer for SPE>=3.
                data_end = (
                    info["xml_footer_offset"]
                    if info["file_header_ver"] >= 3
                    else os.path.getsize(self.request.get_local_filename())
                )
                l = data_end - Spec.data_start
                l //= self._shape[0] * self._shape[1] * self._dtype.itemsize
                if l != self._len:
                    logger.warning(
                        "The file header of %s claims there are %s frames, "
                        "but there are actually %s frames.",
                        self.request.filename,
                        self._len,
                        l,
                    )
                    self._len = min(l, self._len)

            self._meta = None

        def _get_meta_data(self, index):
            if self._meta is None:
                if self._file_header_ver < 3:
                    self._init_meta_data_pre_v3()
                else:
                    self._init_meta_data_post_v3()
            return self._meta

        def _close(self):
            # The file should be closed by `self.request`
            pass

        def _init_meta_data_pre_v3(self):
            self._meta = self._parse_header(Spec.metadata)

            nr = self._meta.pop("NumROI", None)
            nr = 1 if nr < 1 else nr
            self._meta["ROIs"] = roi_array_to_dict(self._meta["ROIs"][:nr])

            # chip sizes
            self._meta["chip_size"] = [
                self._meta.pop("xDimDet", None),
                self._meta.pop("yDimDet", None),
            ]
            self._meta["virt_chip_size"] = [
                self._meta.pop("VChipXdim", None),
                self._meta.pop("VChipYdim", None),
            ]
            self._meta["pre_pixels"] = [
                self._meta.pop("XPrePixels", None),
                self._meta.pop("YPrePixels", None),
            ]
            self._meta["post_pixels"] = [
                self._meta.pop("XPostPixels", None),
                self._meta.pop("YPostPixels", None),
            ]

            # comments
            self._meta["comments"] = [str(c) for c in self._meta["comments"]]

            # geometric operations
            g = []
            f = self._meta.pop("geometric", 0)
            if f & 1:
                g.append("rotate")
            if f & 2:
                g.append("reverse")
            if f & 4:
                g.append("flip")
            self._meta["geometric"] = g

            # Make some additional information more human-readable
            t = self._meta["type"]
            if 1 <= t <= len(Spec.controllers):
                self._meta["type"] = Spec.controllers[t - 1]
            else:
                self._meta["type"] = ""
            m = self._meta["readout_mode"]
            if 1 <= m <= len(Spec.readout_modes):
                self._meta["readout_mode"] = Spec.readout_modes[m - 1]
            else:
                self._meta["readout_mode"] = ""

            # bools
            for k in (
                "absorb_live",
                "can_do_virtual_chip",
                "threshold_min_live",
                "threshold_max_live",
            ):
                self._meta[k] = bool(self._meta[k])

            # frame shape
            self._meta["frame_shape"] = self._shape

            # Extract SDT-control metadata if desired
            if self._sdt_meta:
                SDTControlSpec.extract_metadata(self._meta,
                                                self._char_encoding)

        def _parse_header(self, spec):
            ret = {}
            # Decode each string from the numpy array read by np.fromfile
            decode = np.vectorize(lambda x: x.decode(self._char_encoding))

            for name, sp in spec.items():
                self._file.seek(sp[0])
                cnt = 1 if len(sp) < 3 else sp[2]
                v = np.fromfile(self._file, dtype=sp[1], count=cnt)
                if v.dtype.kind == "S" and name not in Spec.no_decode:
                    # Silently ignore string decoding failures
                    try:
                        v = decode(v)
                    except Exception:
                        logger.warning(
                            'Failed to decode "{}" metadata '
                            "string. Check `char_encoding` "
                            "parameter.".format(name)
                        )

                try:
                    # For convenience, if the array contains only one single
                    # entry, return this entry itself.
                    v = v.item()
                except ValueError:
                    v = np.squeeze(v)
                ret[name] = v
            return ret

        def _init_meta_data_post_v3(self):
            info = self._parse_header(Spec.basic)
            self._file.seek(info["xml_footer_offset"])
            xml = self._file.read()
            self._meta = {"__xml": xml}

        def _get_length(self):
            if self.request.mode[1] in "vV":
                return 1
            else:
                return self._len

        def _get_data(self, index):
            if index < 0:
                raise IndexError("Image index %i < 0" % index)
            if index >= self._len:
                raise IndexError("Image index %i > %i" % (index, self._len))

            if self.request.mode[1] in "vV":
                if index != 0:
                    raise IndexError("Index has to be 0 in v and V modes")
                self._file.seek(Spec.data_start)
                data = np.fromfile(
                    self._file,
                    dtype=self._dtype,
                    count=self._shape[0] * self._shape[1] * self._len,
                )
                data = data.reshape((self._len,) + self._shape)
            else:
                self._file.seek(
                    Spec.data_start
                    + index * self._shape[0] * self._shape[1] * self._dtype.itemsize
                )
                data = np.fromfile(
                    self._file, dtype=self._dtype, count=self._shape[0] * self._shape[1]
                )
                data = data.reshape(self._shape)
            return data, self._get_meta_data(index)


def roi_array_to_dict(a):
    """Convert the `ROIs` structured arrays to :py:class:`dict`

    Parameters
    ----------
    a : numpy.ndarray
        Structured array containing ROI data

    Returns
    -------
    list of dict
        One dict per ROI. Keys are "top_left", "bottom_right", and "bin",
        values are tuples whose first element is the x axis value and the
        second element is the y axis value.
    """
    l = []
    a = a[["startx", "starty", "endx", "endy", "groupx", "groupy"]]
    for sx, sy, ex, ey, gx, gy in a:
        d = {
            "top_left": [int(sx), int(sy)],
            "bottom_right": [int(ex), int(ey)],
            "bin": [int(gx), int(gy)],
        }
        l.append(d)
    return l


fmt = SpeFormat("spe", "SPE file format", ".spe", "iIvV")
formats.add_format(fmt, overwrite=True)
